#pragma once
#include "../CGIResource.h"
#include "VulkanMemory.h"

namespace Alice{
    class VulkanDevice;
    struct CGIVulkanRHIExternalImageDeleteCallbackInfo;
    class VulkanCommandListContext;
    extern VkFormat GVulkanSRGBFormat[EPixelFormat::PF_NUM];
    VkFormat UEToVkTextureFormat(EPixelFormat UEFormat, const bool bIsSRGB);
    VkImageAspectFlags GetAspectMaskFromUEFormat(EPixelFormat Format, bool bIncludeStencil, bool bIncludeDepth = true);
    VkImageViewType UETextureDimensionToVkImageViewType(ETextureDimension Dimension);
	class VulkanResourceMultiBuffer : public CGIBuffer
	{
	public:
		struct FBufferAlloc
		{
			FVulkanAllocation Alloc;
			void* HostPtr = nullptr;
			VkDeviceAddress DeviceAddress = 0;
		};
		VulkanDevice*mDevice;
		VkBufferUsageFlags BufferUsageFlags;
		FBufferAlloc mAlloc;
		VulkanResourceMultiBuffer(VulkanDevice* InDevice, CGIBufferDesc const& InBufferDesc, CGIResourceCreateInfo& CreateInfo, class CGICommandListBase* InRHICmdList = nullptr);
		static VkBufferUsageFlags UEToVKBufferUsageFlags(VulkanDevice* InDevice, EBufferUsageFlags InUEUsage, bool bZeroSize);
		void AllocateMemory(FBufferAlloc& OutAlloc);
		void* Lock(CGICommandListBase& RHICmdList, EResourceLockMode LockMode, uint32 Size, uint32 Offset);
		void Unlock(CGICommandListBase& RHICmdList);
	};
	class CGIVulkanView;
    class CGIVulkanTexture : public CGITexture{
    public:
        VkImage mImage;
	    VkImageUsageFlags mImageUsageFlags;
        VkFormat mStorageFormat;  // Removes SRGB if requested, used to upload data
        VkFormat mViewFormat;  // Format for SRVs, render targets
        VkImageAspectFlags mFullAspectMask;
        VkImageAspectFlags mPartialAspectMask;
	    VkImageTiling mTiling;
	    VkImageLayout mDefaultLayout;
    	CGIVulkanView* mDefaultView;
        CGIVulkanTexture();
        CGIVulkanTexture(VulkanDevice& InDevice, const CGITextureCreateDesc& InCreateDesc, VkImage InImage, const CGIVulkanRHIExternalImageDeleteCallbackInfo& InExternalImageDeleteCallbackInfo);
        // Full includes Depth+Stencil
        inline VkImageAspectFlags GetFullAspectMask() const
        {
            return mFullAspectMask;
        }

        // Only Depth or Stencil
        inline VkImageAspectFlags GetPartialAspectMask() const
        {
            return mPartialAspectMask;
        }
        inline bool IsDepthOrStencilAspect() const
        {
            return (mFullAspectMask & (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT)) != 0;
        }
        inline bool SupportsSampling() const
        {
            return EnumHasAllFlags(GPixelFormats[(uint32)GetDesc().Format].Capabilities, EPixelFormatCapabilities::TextureSample) &&
                ((mImageUsageFlags & VK_IMAGE_USAGE_SAMPLED_BIT) != 0);
        }
        inline VkImageViewType GetViewType() const
        {
            return UETextureDimensionToVkImageViewType(GetDesc().Dimension);
        }
        inline uint32 GetNumberOfArrayLevels() const
        {
            switch (GetViewType())
            {
            case VK_IMAGE_VIEW_TYPE_1D:
            case VK_IMAGE_VIEW_TYPE_2D:
            case VK_IMAGE_VIEW_TYPE_3D:
                return 1;
            case VK_IMAGE_VIEW_TYPE_2D_ARRAY:
                return GetDesc().ArraySize;
            case VK_IMAGE_VIEW_TYPE_CUBE:
                return 6;
            case VK_IMAGE_VIEW_TYPE_CUBE_ARRAY:
                return 6 * GetDesc().ArraySize;
            default:
                return 1;
            }
        }
        void SetInitialImageState(VulkanCommandListContext& Context, VkImageLayout InitialLayout, bool bClear, const FClearValueBinding& ClearValueBinding, bool bIsTransientResource);
    };
    struct VulkanTextureView
    {
        VkImageView View   = VK_NULL_HANDLE;
        VkImage     Image  = VK_NULL_HANDLE;
        uint32      ViewId = 0;
    };
    class CGIVulkanView
    {
    public:
        VulkanDevice& mDevice;
        VulkanTextureView mTextureView;
        CGIVulkanView(VulkanDevice& InDevice, VkDescriptorType InDescriptorType);
        CGIVulkanView* InitAsTextureView(
          VkImage InImage
        , VkImageViewType ViewType
        , VkImageAspectFlags AspectFlags
        , EPixelFormat UEFormat
        , VkFormat Format
        , uint32 FirstMip
        , uint32 NumMips
        , uint32 ArraySliceIndex
        , uint32 NumArraySlices
        , bool bUseIdentitySwizzle = false
        , VkImageUsageFlags ImageUsageFlags = 0
        , VkSamplerYcbcrConversion SamplerYcbcrConversion = nullptr);
		VulkanTextureView const& GetTextureView () const { return mTextureView; }
    };
	
    class VulkanRenderTargetLayout
    {
    public:
        VkAttachmentReference ColorReferences[MaxSimultaneousRenderTargets];
        VkAttachmentReference DepthReference;
		VkAttachmentReferenceStencilLayout StencilReference;

        // Depth goes in the "+1" slot, Depth resolve goes in the "+2 slot", and the Shading Rate texture goes in the "+3" slot.
        VkAttachmentDescription Desc[MaxSimultaneousRenderTargets * 2 + 3];
    	VkAttachmentDescriptionStencilLayout StencilDesc;

        uint8 NumAttachmentDescriptions;
        uint8 NumColorAttachments;
        uint8 bHasDepthStencil;
        uint8 NumUsedClearValues;
		ESubpassHint SubpassHint = ESubpassHint::None;
    	union
    	{
    		VkOffset3D Offset3D;
    		VkOffset2D Offset2D;
    	} Offset;

    	union
    	{
    		VkExtent3D	Extent3D;
    		VkExtent2D	Extent2D;
    	} Extent;
    	VulkanRenderTargetLayout(const FGraphicsPipelineStateInitializer& Initializer);
    	VulkanRenderTargetLayout(VulkanDevice& InDevice, const CGIRenderPassInfo* RPInfo, VkImageLayout CurrentDepthLayout, VkImageLayout CurrentStencilLayout);
    	inline const VkOffset2D& GetOffset2D() const { return Offset.Offset2D; }
    	inline const VkOffset3D& GetOffset3D() const { return Offset.Offset3D; }
    	inline const VkExtent2D& GetExtent2D() const { return Extent.Extent2D; }
    	inline const VkExtent3D& GetExtent3D() const { return Extent.Extent3D; }
		inline const ESubpassHint GetSubpassHint() const { return SubpassHint; }
    	inline uint32 GetNumAttachmentDescriptions() const { return NumAttachmentDescriptions; }
    	inline const VkAttachmentDescription* GetAttachmentDescriptions() const { return Desc; }
    	
    	inline const VkAttachmentReference* GetColorAttachmentReferences() const { return NumColorAttachments > 0 ? ColorReferences : nullptr; }
    	inline const VkAttachmentReference* GetDepthAttachmentReference() const { return bHasDepthStencil ? &DepthReference : nullptr; }
    	inline const VkAttachmentReferenceStencilLayout* GetStencilAttachmentReference() const { return bHasDepthStencil ? &StencilReference : nullptr; }
    	inline const VkAttachmentDescriptionStencilLayout* GetStencilDesc() const { return bHasDepthStencil ? &StencilDesc : nullptr; }
    	bool operator==(const VulkanRenderTargetLayout&inRenderTargetLayout)const;
    };
	class VulkanRenderPass
	{
	public:
		inline const VulkanRenderTargetLayout& GetLayout() const
		{
			return mLayout;
		}

		inline VkRenderPass GetHandle() const
		{
			return mRenderPass;
		}

		inline uint32 GetNumUsedClearValues() const
		{
			return mNumUsedClearValues;
		}

		friend class FVulkanRenderPassManager;
		friend class FVulkanPipelineStateCacheManager;

		VulkanRenderPass(VulkanDevice& Device, const VulkanRenderTargetLayout& RTLayout);
		~VulkanRenderPass();

		VulkanRenderTargetLayout	mLayout;
		VkRenderPass				mRenderPass;
		uint32						mNumUsedClearValues;
		VulkanDevice&			mDevice;
	};
	class VulkanFramebuffer
	{
	public:
		const VulkanRenderPass* mRenderPass;
		VkImageView mColorAttachment0ImageView;
		VulkanFramebuffer(VulkanDevice& Device, const CGISetRenderTargetsInfo& InRTInfo, const VulkanRenderTargetLayout& RTLayout, const VulkanRenderPass& RenderPass);
		~VulkanFramebuffer();

		bool Matches(const CGISetRenderTargetsInfo& RTInfo) const;

		uint32 GetNumColorAttachments() const
		{
			return NumColorAttachments;
		}

		void Destroy(VulkanDevice& Device);

		VkFramebuffer GetHandle()
		{
			return Framebuffer;
		}

		bool ContainsRenderTarget(CGITexture* Texture) const
		{
			CGIVulkanTexture* VulkanTexture = reinterpret_cast<CGIVulkanTexture*>(Texture);
			return ContainsRenderTarget(VulkanTexture->mImage);
		}

		bool ContainsRenderTarget(VkImage Image) const
		{
			for (uint32 Index = 0; Index < NumColorAttachments; ++Index)
			{
				if (ColorRenderTargetImages[Index] == Image)
				{
					return true;
				}
			}

			return (DepthStencilRenderTargetImage == Image);
		}

		VkRect2D GetRenderArea() const
		{
			return RenderArea;
		}

	private:
		VkFramebuffer Framebuffer;
		VkRect2D RenderArea;

		// Unadjusted number of color render targets as in FRHISetRenderTargetsInfo 
		uint32 NumColorRenderTargets;

		// Save image off for comparison, in case it gets aliased.
		uint32 NumColorAttachments;
		VkImage ColorRenderTargetImages[MaxSimultaneousRenderTargets];
		VkImage ColorResolveTargetImages[MaxSimultaneousRenderTargets];
		VkImage DepthStencilRenderTargetImage;
		VkImage DepthStencilResolveRenderTargetImage;
		VkImage FragmentDensityImage;

		// Predefined set of barriers, when executes ensuring all writes are finished
		std::vector<VkImageMemoryBarrier> WriteBarriers;
		std::vector<CGIVulkanView> OwnedTextureViews;
		std::vector<CGIVulkanView const*> AttachmentTextureViews;
		friend class FVulkanCommandListContext;
	};
}